{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "415e3d92-7369-45fe-a7ea-b0c982176f38",
   "metadata": {},
   "source": [
    "# SplitRec：在隐语拆分学习中使用 LabelInferenceAttack\n",
    "在联邦学习中，攻击者可以通过监听训练模型过程中传输的数值和梯度信息，攻击对方模型或数据，在一定程度上推理出有用信息，造成信息泄露。\n",
    "\n",
    "本文考虑两方拆分学习中的标签推理攻击，将介绍论文[《Label Inference Attacks Against Vertical Federated Learning》](https://www.usenix.org/system/files/sec22summer_fu.pdf)中的攻击方法在隐语中的使用。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95dc8d90-4e12-464b-9c7d-0c9901197f4b",
   "metadata": {},
   "source": [
    "## Label Inference Attack Against Vertical Federated Learning\n",
    "论文中提出纵向联邦场景下主动和被动两种标签推理攻击方法，攻击方推测对方的样本标签造成信息泄露。\n",
    "\n",
    "![lia0](./resources/lia0.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56ac3058-e1db-427a-8aff-37e3b587bb5c",
   "metadata": {},
   "source": [
    "### Passive Label Inference Attack through Model Completion\n",
    "被动的标签推理攻击发生在联邦模型训练之后，攻击模型推断对方标签，联邦模型训练期间不会受到干扰。\n",
    "\n",
    "攻击模型在本地 bottom 模型（SL Model 中的 base 模型）基础上增加 top 模型使其可以预测标签。\n",
    "\n",
    "攻击模型利用少量带标签的辅助数据和大量无标签数据，使用改进的半监督学习算法 MixMatch 进行训练。MixMatch 通过对少量标记数据和未标记数据进行数据增广，分别计算两者的损失，整体损失是两者的加权。具体改进的 MixMatch 算法如下\n",
    "\n",
    "![lia1](./resources/lia1.png)  \n",
    "![lia2](./resources/lia2.png)\n",
    "\n",
    "其中 loss 的计算公式如下\n",
    "\n",
    "\n",
    "![lia3](./resources/lia3.png)  \n",
    "\n",
    "![lia4](./resources/lia4.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5282d99-7e7b-48b3-b9f9-ed9e9a632f63",
   "metadata": {},
   "source": [
    "### Active Label Inference Attack with the Malicious Local Optimizer\n",
    "在 Passive Label Inference Attack 基础上，恶意的攻击者可以通过使用自适应优化器扩大学习率，使 top 模型更多依赖自己的 bottom 模型，进而提高攻击的效果。\n",
    "\n",
    "恶意优化算法如下\n",
    "\n",
    "![lia5](./resources/lia5.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f3b1612-3dd8-45d9-956d-040e3df12bd3",
   "metadata": {},
   "source": [
    "## 隐语中的攻击方法实现\n",
    "在隐语中攻击方法的实现是通过 callback 机制来完成。攻击算法基类 CallBack 位于 secretflow/ml/nn/sl/backend/torch/callback.py，我们在联邦模型训练的以下几个节点提供 hook，不同攻击方法可以通过将攻击算法实现在对应节点的 hook， 使攻击逻辑注入到联邦模型的训练过程中。\n",
    "\n",
    "- on_train_begin\n",
    "- on_train_end\n",
    "- on_epoch_begin\n",
    "- on_epoch_end\n",
    "- on_batch_begin\n",
    "- on_batch_end\n",
    "\n",
    "用户如果需要实现自定义的攻击方法，需要\n",
    "\n",
    "1. 定义 CustomAttacker 继承基类 Callback，将攻击逻辑实现到对应的 hook 函数中\n",
    "2. 定义 attacker_builder 函数将构建 attacker 写到其中\n",
    "3. 与普通 Split Learning 模型训练一样定义 sl_model, 并在调用 sl_model.fit() 时，将 callback_dict {party -> attacker_builder} 传入 callbacks 参数即可\n",
    "\n",
    "其中步骤 1 可以参考隐语中已有的 FeatureInferenceAttacker/LabelInferenceAttacker，步骤 2 和 3 可参考下面 LeatureInferenceAttacker 的使用方式。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d6db125-6aee-45bf-99b4-6eb211375bfe",
   "metadata": {},
   "source": [
    "## Label Inferece Attack 的隐语封装\n",
    "我们在隐语中提供了多种攻击方法的封装。对于论文中的攻击方法，我们提供了 LabelInferenceAttacker 封装，具体使用可以参考以下代码。\n",
    "\n",
    "首先和一般 Split Learning 模型训练一样，我们将进行数据处理，并定义一个 SLModel。\n",
    "\n",
    "然后定义调用 LeatureInferenceAttacker 的 attacker_builder，并在 SLModel fit 时将 attacker_builder 传入进行训练和攻击。 "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b3fe478-e3e5-4b07-87d1-8234b6eff490",
   "metadata": {},
   "source": [
    "## 环境设置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b67cf5cb-498d-42f6-9c47-14a6a5af407c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The version of SecretFlow: 1.1.0.dev20230926\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-09-26 19:49:33,980\tINFO worker.py:1538 -- Started a local Ray instance.\n"
     ]
    }
   ],
   "source": [
    "import secretflow as sf\n",
    "\n",
    "# Check the version of your SecretFlow\n",
    "print('The version of SecretFlow: {}'.format(sf.__version__))\n",
    "\n",
    "# In case you have a running secretflow runtime already.\n",
    "sf.shutdown()\n",
    "sf.init(['alice', 'bob'], address=\"local\")\n",
    "alice, bob = sf.PYU('alice'), sf.PYU('bob')\n",
    "device_y = bob"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02c1ef3a-4d00-46da-bcc5-8b2884c7eca0",
   "metadata": {},
   "source": [
    "## 数据集介绍\n",
    "\n",
    "这里我们使用经典图像数据集 Cifar10。\n",
    "\n",
    "这里我们对数据进行拆分，拆分模型两方各持有半张图片，攻击方不持有标签。\n",
    "\n",
    "[数据集官网](https://www.cs.toronto.edu/~kriz/cifar.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a3d9f1d-9621-4cb6-b51d-2b4b2d1e101a",
   "metadata": {},
   "source": [
    "## 准备训练数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4c7ac7e4-4c8d-415c-aa65-b9a13fbd44cd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz to ./data_download/cifar-10-python.tar.gz\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 170498071/170498071 [15:08<00:00, 187619.24it/s] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting ./data_download/cifar-10-python.tar.gz to ./data_download\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-2-8b71ce40e487>:11: FutureWarning: The input object of type 'Tensor' is an array-like implementing one of the corresponding protocols (`__array__`, `__array_interface__` or `__array_struct__`); but not a sequence (or 0-D). In the future, this object will be coerced as if it was first converted using `np.array(obj)`. To retain the old behaviour, you have to either modify the type 'Tensor', or assign to an empty array created with `np.empty(correct_shape, dtype=object)`.\n",
      "  train_np = np.array(train_loader.dataset)\n",
      "<ipython-input-2-8b71ce40e487>:11: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray.\n",
      "  train_np = np.array(train_loader.dataset)\n"
     ]
    }
   ],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets, transforms\n",
    "import numpy as np\n",
    "from secretflow.data.ndarray import FedNdarray, PartitionWay\n",
    "\n",
    "data_file_path = './data_download'\n",
    "batch_size = 128\n",
    "train_dataset = datasets.CIFAR10(\n",
    "    data_file_path, True, transform=transforms.ToTensor(), download=True\n",
    ")\n",
    "train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "train_np = np.array(train_loader.dataset)\n",
    "train_data = np.array([t[0].numpy() for t in train_np])\n",
    "train_label = np.array([t[1] for t in train_np])\n",
    "\n",
    "fed_data = FedNdarray(\n",
    "    partitions={\n",
    "        alice: alice(lambda x: x[:, :, :, 0:16])(train_data),\n",
    "        bob: bob(lambda x: x[:, :, :, 16:32])(train_data),\n",
    "    },\n",
    "    partition_way=PartitionWay.VERTICAL,\n",
    ")\n",
    "# party bob hold label\n",
    "label = bob(lambda x: x)(train_label)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "727d8377-959c-4a15-8a34-cc4e17eb465b",
   "metadata": {},
   "source": [
    "## 定义SLModel的模型结构"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2af1c2f2-f23a-4675-ae92-f2ec9bdf57ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.nn.init as init\n",
    "\n",
    "\n",
    "def weights_init(m):\n",
    "    if isinstance(m, nn.Linear) or isinstance(m, nn.Conv2d):\n",
    "        init.kaiming_normal_(m.weight)\n",
    "\n",
    "\n",
    "class LambdaLayer(nn.Module):\n",
    "    def __init__(self, lambd):\n",
    "        super(LambdaLayer, self).__init__()\n",
    "        self.lambd = lambd\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.lambd(x)\n",
    "\n",
    "\n",
    "class BasicBlock(nn.Module):\n",
    "    expansion = 1\n",
    "\n",
    "    def __init__(self, in_planes, planes, kernel_size, stride=1, option='A'):\n",
    "        super(BasicBlock, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(\n",
    "            in_planes,\n",
    "            planes,\n",
    "            kernel_size=kernel_size,\n",
    "            stride=stride,\n",
    "            padding=1,\n",
    "            bias=False,\n",
    "        )\n",
    "        self.bn1 = nn.BatchNorm2d(planes)\n",
    "        self.conv2 = nn.Conv2d(\n",
    "            planes, planes, kernel_size=kernel_size, stride=1, padding=1, bias=False\n",
    "        )\n",
    "        self.bn2 = nn.BatchNorm2d(planes)\n",
    "\n",
    "        self.shortcut = nn.Sequential()\n",
    "        if stride != 1 or in_planes != planes:\n",
    "            if option == 'A':\n",
    "                \"\"\"\n",
    "                For CIFAR10 ResNet paper uses option A.\n",
    "                \"\"\"\n",
    "                self.shortcut = LambdaLayer(\n",
    "                    lambda x: F.pad(\n",
    "                        x[:, :, ::2, ::2],\n",
    "                        (0, 0, 0, 0, planes // 4, planes // 4),\n",
    "                        \"constant\",\n",
    "                        0,\n",
    "                    )\n",
    "                )\n",
    "            elif option == 'B':\n",
    "                self.shortcut = nn.Sequential(\n",
    "                    nn.Conv2d(\n",
    "                        in_planes,\n",
    "                        self.expansion * planes,\n",
    "                        kernel_size=1,\n",
    "                        stride=stride,\n",
    "                        bias=False,\n",
    "                    ),\n",
    "                    nn.BatchNorm2d(self.expansion * planes),\n",
    "                )\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = F.relu(self.bn1(self.conv1(x)))\n",
    "        out = self.bn2(self.conv2(out))\n",
    "        out += self.shortcut(x)\n",
    "        out = F.relu(out)\n",
    "        return out\n",
    "\n",
    "\n",
    "class ResNet(nn.Module):\n",
    "    def __init__(self, block, num_blocks, kernel_size, num_classes):\n",
    "        super(ResNet, self).__init__()\n",
    "        self.in_planes = 16\n",
    "\n",
    "        self.conv1 = nn.Conv2d(\n",
    "            3, 16, kernel_size=kernel_size, stride=1, padding=1, bias=False\n",
    "        )\n",
    "        self.bn1 = nn.BatchNorm2d(16)\n",
    "        self.layer1 = self._make_layer(block, 16, num_blocks[0], kernel_size, stride=1)\n",
    "        self.layer2 = self._make_layer(block, 32, num_blocks[1], kernel_size, stride=2)\n",
    "        self.layer3 = self._make_layer(block, 64, num_blocks[2], kernel_size, stride=2)\n",
    "        self.linear = nn.Linear(64, num_classes, bias=False)\n",
    "\n",
    "        self.apply(weights_init)\n",
    "\n",
    "    def _make_layer(self, block, planes, num_blocks, kernel_size, stride):\n",
    "        strides = [stride] + [1] * (num_blocks - 1)\n",
    "        layers = []\n",
    "        for stride in strides:\n",
    "            layers.append(block(self.in_planes, planes, kernel_size, stride))\n",
    "            self.in_planes = planes * block.expansion\n",
    "\n",
    "        return nn.Sequential(*layers)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # [bs,3,32,16]\n",
    "        out = F.relu(self.bn1(self.conv1(x)))\n",
    "        # [bs,16,32,16]\n",
    "        out = self.layer1(out)\n",
    "        # [bs,16,32,16]\n",
    "        out = self.layer2(out)\n",
    "        # [bs,32,16,8]\n",
    "        out = self.layer3(out)\n",
    "        # [bs,64,8,4]\n",
    "        out = F.avg_pool2d(out, out.size()[2:])\n",
    "        # [bs,64,1,1]\n",
    "        out = out.view(out.size(0), -1)\n",
    "        # [bs,64]\n",
    "        out = self.linear(out)\n",
    "        # [bs,10]\n",
    "        return out\n",
    "\n",
    "\n",
    "def resnet20(kernel_size=(3, 3), num_classes=10):\n",
    "    return ResNet(\n",
    "        block=BasicBlock,\n",
    "        num_blocks=[3, 3, 3],\n",
    "        kernel_size=kernel_size,\n",
    "        num_classes=num_classes,\n",
    "    )\n",
    "\n",
    "\n",
    "# base model\n",
    "class BottomModelForCifar10(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(BottomModelForCifar10, self).__init__()\n",
    "        self.resnet20 = resnet20(num_classes=10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.resnet20(x)\n",
    "        return x\n",
    "\n",
    "    def output_num(self):\n",
    "        return 1\n",
    "\n",
    "\n",
    "# fuse model\n",
    "class TopModelForCifar10(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(TopModelForCifar10, self).__init__()\n",
    "        self.fc1top = nn.Linear(20, 20)\n",
    "        self.fc2top = nn.Linear(20, 10)\n",
    "        self.fc3top = nn.Linear(10, 10)\n",
    "        self.fc4top = nn.Linear(10, 10)\n",
    "        self.bn0top = nn.BatchNorm1d(20)\n",
    "        self.bn1top = nn.BatchNorm1d(20)\n",
    "        self.bn2top = nn.BatchNorm1d(10)\n",
    "        self.bn3top = nn.BatchNorm1d(10)\n",
    "        print('batch norm: ', self.bn0top)\n",
    "        self.apply(weights_init)\n",
    "\n",
    "    def forward(self, input_tensor):\n",
    "        output_bottom_models = torch.cat(input_tensor, dim=1)\n",
    "        x = output_bottom_models\n",
    "        x = self.fc1top(F.relu(self.bn0top(x)))\n",
    "        x = self.bn1top(x)\n",
    "        x = self.fc2top(F.relu(x))\n",
    "        x = self.fc3top(F.relu(self.bn2top(x)))\n",
    "        x = self.fc4top(F.relu(self.bn3top(x)))\n",
    "        return F.log_softmax(x, dim=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a529bdf-71a1-4a4a-9244-9d0abce3fd83",
   "metadata": {},
   "source": [
    "## 定义SLModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ca79c923-0699-43bc-9f60-63af671c53c7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Create proxy actor <class 'secretflow_fl.ml.nn.sl.backend.torch.strategy.split_nn.PYUSLTorchModel'> with party alice.\n",
      "INFO:root:Create proxy actor <class 'secretflow_fl.ml.nn.sl.backend.torch.strategy.split_nn.PYUSLTorchModel'> with party bob.\n"
     ]
    }
   ],
   "source": [
    "import torch.optim as optim\n",
    "from torchmetrics import Accuracy, Precision\n",
    "from secretflow_fl.ml.nn import SLModel\n",
    "from secretflow_fl.ml.nn.core.torch import metric_wrapper, optim_wrapper, TorchModel\n",
    "\n",
    "\n",
    "loss_fn = nn.CrossEntropyLoss\n",
    "optim_fn = optim_wrapper(optim.SGD, lr=1e-2, momentum=0.9, weight_decay=5e-4)\n",
    "base_model = TorchModel(\n",
    "    model_fn=BottomModelForCifar10,\n",
    "    loss_fn=loss_fn,\n",
    "    optim_fn=optim_fn,\n",
    "    metrics=[\n",
    "        metric_wrapper(Accuracy, task=\"multiclass\", num_classes=10, average='micro'),\n",
    "        metric_wrapper(Precision, task=\"multiclass\", num_classes=10, average='micro'),\n",
    "    ],\n",
    ")\n",
    "\n",
    "fuse_model = TorchModel(\n",
    "    model_fn=TopModelForCifar10,\n",
    "    loss_fn=loss_fn,\n",
    "    optim_fn=optim_fn,\n",
    "    metrics=[\n",
    "        metric_wrapper(Accuracy, task=\"multiclass\", num_classes=10, average='micro'),\n",
    "        metric_wrapper(Precision, task=\"multiclass\", num_classes=10, average='micro'),\n",
    "    ],\n",
    ")\n",
    "\n",
    "base_model_dict = {\n",
    "    alice: base_model,\n",
    "    bob: base_model,\n",
    "}\n",
    "\n",
    "sl_model = SLModel(\n",
    "    base_model_dict=base_model_dict,\n",
    "    device_y=device_y,\n",
    "    model_fuse=fuse_model,\n",
    "    dp_strategy_dict=None,\n",
    "    compressor=None,\n",
    "    simulation=True,\n",
    "    random_seed=1234,\n",
    "    backend='torch',\n",
    "    strategy='split_nn',\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78b66125-6e5c-4dc1-82e1-39146e804a01",
   "metadata": {},
   "source": [
    "## 定义 attacker_builder\n",
    "### 定义 Attacker 模型结构\n",
    "这里 Attacker 模型是在原 base 模型（`BottomModelForCifar10`）基础上加一个 top 模型，这里只需定义 top 模型 `BottomModelPlus`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3acb9532-7ae0-475b-981c-17c41e1d78d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "def weights_init_ones(m):\n",
    "    if isinstance(m, nn.Linear) or isinstance(m, nn.Conv2d):\n",
    "        init.ones_(m.weight)\n",
    "\n",
    "\n",
    "# for attacker\n",
    "class BottomModelPlus(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        bottom_model,\n",
    "        size_bottom_out=10,\n",
    "        num_classes=10,\n",
    "        num_layer=1,\n",
    "        activation_func_type='ReLU',\n",
    "        use_bn=True,\n",
    "    ):\n",
    "        super(BottomModelPlus, self).__init__()\n",
    "        self.bottom_model = bottom_model\n",
    "\n",
    "        dict_activation_func_type = {'ReLU': F.relu, 'Sigmoid': F.sigmoid, 'None': None}\n",
    "        self.activation_func = dict_activation_func_type[activation_func_type]\n",
    "        self.num_layer = num_layer\n",
    "        self.use_bn = use_bn\n",
    "\n",
    "        self.fc_1 = nn.Linear(size_bottom_out, size_bottom_out, bias=True)\n",
    "        self.bn_1 = nn.BatchNorm1d(size_bottom_out)\n",
    "        self.fc_1.apply(weights_init_ones)\n",
    "\n",
    "        self.fc_2 = nn.Linear(size_bottom_out, size_bottom_out, bias=True)\n",
    "        self.bn_2 = nn.BatchNorm1d(size_bottom_out)\n",
    "        self.fc_2.apply(weights_init_ones)\n",
    "\n",
    "        self.fc_3 = nn.Linear(size_bottom_out, size_bottom_out, bias=True)\n",
    "        self.bn_3 = nn.BatchNorm1d(size_bottom_out)\n",
    "        self.fc_3.apply(weights_init_ones)\n",
    "\n",
    "        self.fc_4 = nn.Linear(size_bottom_out, size_bottom_out, bias=True)\n",
    "        self.bn_4 = nn.BatchNorm1d(size_bottom_out)\n",
    "        self.fc_4.apply(weights_init_ones)\n",
    "\n",
    "        self.fc_final = nn.Linear(size_bottom_out, num_classes, bias=True)\n",
    "        self.bn_final = nn.BatchNorm1d(size_bottom_out)\n",
    "        self.fc_final.apply(weights_init_ones)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.bottom_model(x)\n",
    "\n",
    "        if self.num_layer >= 2:\n",
    "            if self.use_bn:\n",
    "                x = self.bn_1(x)\n",
    "            if self.activation_func:\n",
    "                x = self.activation_func(x)\n",
    "            x = self.fc_1(x)\n",
    "\n",
    "        if self.num_layer >= 3:\n",
    "            if self.use_bn:\n",
    "                x = self.bn_2(x)\n",
    "            if self.activation_func:\n",
    "                x = self.activation_func(x)\n",
    "            x = self.fc_2(x)\n",
    "\n",
    "        if self.num_layer >= 4:\n",
    "            if self.use_bn:\n",
    "                x = self.bn_3(x)\n",
    "            if self.activation_func:\n",
    "                x = self.activation_func(x)\n",
    "            x = self.fc_3(x)\n",
    "\n",
    "        if self.num_layer >= 5:\n",
    "            if self.use_bn:\n",
    "                x = self.bn_4(x)\n",
    "            if self.activation_func:\n",
    "                x = self.activation_func(x)\n",
    "            x = self.fc_4(x)\n",
    "        if self.use_bn:\n",
    "            x = self.bn_final(x)\n",
    "        if self.activation_func:\n",
    "            x = self.activation_func(x)\n",
    "        x = self.fc_final(x)\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3348a444-1e47-42cc-92c7-84bf2e0f30a7",
   "metadata": {},
   "source": [
    "### 定义 Attacker 需要的数据\n",
    "这里标签攻击算法需要少量有标签的辅助数据集，及本身的大量无标签数据。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4a0c2b63-b8f2-461f-a173-434165e07b6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "\n",
    "class CIFAR10Labeled(datasets.CIFAR10):\n",
    "    def __init__(\n",
    "        self,\n",
    "        root,\n",
    "        indexs=None,\n",
    "        train=True,\n",
    "        transform=None,\n",
    "        target_transform=None,\n",
    "        download=False,\n",
    "    ):\n",
    "        super(CIFAR10Labeled, self).__init__(\n",
    "            root,\n",
    "            train=train,\n",
    "            transform=transform,\n",
    "            target_transform=target_transform,\n",
    "            download=download,\n",
    "        )\n",
    "        if indexs is not None:\n",
    "            self.data = self.data[indexs]\n",
    "            self.targets = np.array(self.targets)[indexs]\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            index (int): Index\n",
    "\n",
    "        Returns:\n",
    "            tuple: (image, target) where target is index of the target class.\n",
    "        \"\"\"\n",
    "        img, target = self.data[index], self.targets[index]\n",
    "\n",
    "        if self.transform is not None:\n",
    "            img = self.transform(img)\n",
    "        img = img[:, :, :16]\n",
    "\n",
    "        if self.target_transform is not None:\n",
    "            target = self.target_transform(target)\n",
    "\n",
    "        return img, target\n",
    "\n",
    "\n",
    "class CIFAR10Unlabeled(CIFAR10Labeled):\n",
    "    def __init__(\n",
    "        self,\n",
    "        root,\n",
    "        indexs,\n",
    "        train=True,\n",
    "        transform=None,\n",
    "        target_transform=None,\n",
    "        download=False,\n",
    "    ):\n",
    "        super(CIFAR10Unlabeled, self).__init__(\n",
    "            root,\n",
    "            indexs,\n",
    "            train=train,\n",
    "            transform=transform,\n",
    "            target_transform=target_transform,\n",
    "            download=download,\n",
    "        )\n",
    "        self.targets = np.array([-1 for i in range(len(self.targets))])\n",
    "        self.data = self.data[:, :, :, :16]\n",
    "\n",
    "\n",
    "def label_index_split(labels, n_labeled_per_class, num_classes):\n",
    "    labels = np.array(labels)\n",
    "    train_labeled_idxs = []\n",
    "    train_unlabeled_idxs = []\n",
    "\n",
    "    for i in range(num_classes):\n",
    "        idxs = np.where(labels == i)[0]\n",
    "        np.random.shuffle(idxs)\n",
    "        train_labeled_idxs.extend(idxs[:n_labeled_per_class])\n",
    "        train_unlabeled_idxs.extend(idxs[n_labeled_per_class:])\n",
    "    np.random.shuffle(train_labeled_idxs)\n",
    "    np.random.shuffle(train_unlabeled_idxs)\n",
    "    return train_labeled_idxs, train_unlabeled_idxs\n",
    "\n",
    "\n",
    "def data_builder(batch_size, file_path=None):\n",
    "    def prepare_data():\n",
    "        n_labeled = 40\n",
    "        num_classes = 10\n",
    "\n",
    "        def get_transforms():\n",
    "            transform_ = transforms.Compose(\n",
    "                [\n",
    "                    transforms.ToTensor(),\n",
    "                ]\n",
    "            )\n",
    "            return transform_\n",
    "\n",
    "        transforms_ = get_transforms()\n",
    "\n",
    "        base_dataset = datasets.CIFAR10(file_path, train=True)\n",
    "\n",
    "        train_labeled_idxs, train_unlabeled_idxs = label_index_split(\n",
    "            base_dataset.targets, int(n_labeled / num_classes), num_classes\n",
    "        )\n",
    "        train_labeled_dataset = CIFAR10Labeled(\n",
    "            file_path, train_labeled_idxs, train=True, transform=transforms_\n",
    "        )\n",
    "        train_unlabeled_dataset = CIFAR10Unlabeled(\n",
    "            file_path, train_unlabeled_idxs, train=True, transform=transforms_\n",
    "        )\n",
    "        train_complete_dataset = CIFAR10Labeled(\n",
    "            file_path, None, train=True, transform=transforms_\n",
    "        )\n",
    "        test_dataset = CIFAR10Labeled(\n",
    "            file_path, train=False, transform=transforms_, download=True\n",
    "        )\n",
    "        print(\n",
    "            \"#Labeled:\",\n",
    "            len(train_labeled_idxs),\n",
    "            \"#Unlabeled:\",\n",
    "            len(train_unlabeled_idxs),\n",
    "        )\n",
    "\n",
    "        # auxiliary dataset\n",
    "        labeled_trainloader = torch.utils.data.DataLoader(\n",
    "            train_labeled_dataset,\n",
    "            batch_size=batch_size,\n",
    "            shuffle=True,\n",
    "            num_workers=0,\n",
    "            drop_last=True,\n",
    "        )\n",
    "        unlabeled_trainloader = torch.utils.data.DataLoader(\n",
    "            train_unlabeled_dataset,\n",
    "            batch_size=batch_size,\n",
    "            shuffle=False,\n",
    "            num_workers=0,\n",
    "            drop_last=True,\n",
    "        )\n",
    "        dataset_bs = batch_size * 10\n",
    "        test_loader = torch.utils.data.DataLoader(\n",
    "            test_dataset, batch_size=dataset_bs, shuffle=False, num_workers=0\n",
    "        )\n",
    "        train_complete_trainloader = torch.utils.data.DataLoader(\n",
    "            train_complete_dataset,\n",
    "            batch_size=dataset_bs,\n",
    "            shuffle=False,\n",
    "            num_workers=0,\n",
    "            drop_last=True,\n",
    "        )\n",
    "        return (\n",
    "            labeled_trainloader,\n",
    "            unlabeled_trainloader,\n",
    "            test_loader,\n",
    "            train_complete_trainloader,\n",
    "        )\n",
    "\n",
    "    return prepare_data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c163d4eb-3082-4149-90fd-03a1fd283407",
   "metadata": {},
   "source": [
    "### 定义 attacker_builder\n",
    "这里 attacker_builder 是一个字典，其元素是参与方和对应的 attacker_builder，通常只需要填充攻击方和对应的 attacker_builder_function。\n",
    "\n",
    "本文中 Label Inference Attack 只需填充攻击方 alice 和对应的 attacker_builder_function。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2f30ecd5-9aa7-470a-90ac-d5c0478382c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from secretflow_fl.ml.nn.sl.attacks.lia_torch import (\n",
    "    LabelInferenceAttacker,\n",
    ")\n",
    "\n",
    "\n",
    "def create_attacker_builder(file_path, batch_size, model_save_path):\n",
    "    def attacker_builder():\n",
    "        def create_model(ema=False):\n",
    "            bottom_model = BottomModelForCifar10()\n",
    "            model = BottomModelPlus(bottom_model)\n",
    "\n",
    "            if ema:\n",
    "                for param in model.parameters():\n",
    "                    param.detach_()\n",
    "\n",
    "            return model\n",
    "\n",
    "        model = create_model(ema=False)\n",
    "        ema_model = create_model(ema=True)\n",
    "\n",
    "        data_buil = data_builder(batch_size=batch_size, file_path=file_path)\n",
    "\n",
    "        attacker = LabelInferenceAttacker(\n",
    "            model, ema_model, 10, data_buil, save_model_path=model_save_path\n",
    "        )\n",
    "        return attacker\n",
    "\n",
    "    return attacker_builder\n",
    "\n",
    "\n",
    "callback_dict = {\n",
    "    alice: create_attacker_builder(\n",
    "        file_path=data_file_path, batch_size=16, model_save_path=None\n",
    "    )\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6c33852-d09c-4adb-9976-351031e2d5d8",
   "metadata": {},
   "source": [
    "## 开始训练和攻击"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2f9ec582-00e9-45e5-8fb3-d9359e094409",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:SL Train Params: {'x': FedNdarray(partitions={PYURuntime(alice): <secretflow.device.device.pyu.PYUObject object at 0x7f3728683280>, PYURuntime(bob): <secretflow.device.device.pyu.PYUObject object at 0x7f372336a250>}, partition_way=<PartitionWay.VERTICAL: 'vertical'>), 'y': <secretflow.device.device.pyu.PYUObject object at 0x7f372336a8b0>, 'batch_size': 128, 'epochs': 1, 'verbose': 1, 'callbacks': {PYURuntime(alice): <function create_attacker_builder.<locals>.attacker_builder at 0x7f36ca5eb9d0>}, 'validation_data': (FedNdarray(partitions={PYURuntime(alice): <secretflow.device.device.pyu.PYUObject object at 0x7f3728683280>, PYURuntime(bob): <secretflow.device.device.pyu.PYUObject object at 0x7f372336a250>}, partition_way=<PartitionWay.VERTICAL: 'vertical'>), <secretflow.device.device.pyu.PYUObject object at 0x7f372336a8b0>), 'shuffle': False, 'sample_weight': None, 'validation_freq': 1, 'dp_spent_step_freq': None, 'dataset_builder': None, 'audit_log_params': {}, 'random_seed': 1234, 'audit_log_dir': None, 'self': <secretflow_fl.ml.nn.sl.sl_model.SLModel object at 0x7f3712ff94c0>}\n",
      "\u001b[2m\u001b[36m(pid=1866347)\u001b[0m 2023-09-26 20:05:02.902472: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/rh/gcc-toolset-11/root/usr/lib64:/opt/rh/gcc-toolset-11/root/usr/lib:/opt/rh/gcc-toolset-11/root/usr/lib64/dyninst:/opt/rh/gcc-toolset-11/root/usr/lib/dyninst\n",
      "\u001b[2m\u001b[36m(pid=1866419)\u001b[0m 2023-09-26 20:05:03.037969: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/rh/gcc-toolset-11/root/usr/lib64:/opt/rh/gcc-toolset-11/root/usr/lib:/opt/rh/gcc-toolset-11/root/usr/lib64/dyninst:/opt/rh/gcc-toolset-11/root/usr/lib/dyninst\n",
      "\u001b[2m\u001b[36m(pid=1866347)\u001b[0m 2023-09-26 20:05:03.744462: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/rh/gcc-toolset-11/root/usr/lib64:/opt/rh/gcc-toolset-11/root/usr/lib:/opt/rh/gcc-toolset-11/root/usr/lib64/dyninst:/opt/rh/gcc-toolset-11/root/usr/lib/dyninst\n",
      "\u001b[2m\u001b[36m(pid=1866347)\u001b[0m 2023-09-26 20:05:03.744559: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/rh/gcc-toolset-11/root/usr/lib64:/opt/rh/gcc-toolset-11/root/usr/lib:/opt/rh/gcc-toolset-11/root/usr/lib64/dyninst:/opt/rh/gcc-toolset-11/root/usr/lib/dyninst\n",
      "\u001b[2m\u001b[36m(pid=1866347)\u001b[0m 2023-09-26 20:05:03.744571: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.\n",
      "\u001b[2m\u001b[36m(pid=1866419)\u001b[0m 2023-09-26 20:05:03.855912: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/rh/gcc-toolset-11/root/usr/lib64:/opt/rh/gcc-toolset-11/root/usr/lib:/opt/rh/gcc-toolset-11/root/usr/lib64/dyninst:/opt/rh/gcc-toolset-11/root/usr/lib/dyninst\n",
      "\u001b[2m\u001b[36m(pid=1866419)\u001b[0m 2023-09-26 20:05:03.856003: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/rh/gcc-toolset-11/root/usr/lib64:/opt/rh/gcc-toolset-11/root/usr/lib:/opt/rh/gcc-toolset-11/root/usr/lib64/dyninst:/opt/rh/gcc-toolset-11/root/usr/lib/dyninst\n",
      "\u001b[2m\u001b[36m(pid=1866419)\u001b[0m 2023-09-26 20:05:03.856014: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866419)\u001b[0m batch norm:  BatchNorm1d(20, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|          | 2/391 [00:00<00:45,  8.53it/s]\u001b[2m\u001b[36m(_run pid=1826653)\u001b[0m 2023-09-26 20:05:07.309229: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/rh/gcc-toolset-11/root/usr/lib64:/opt/rh/gcc-toolset-11/root/usr/lib:/opt/rh/gcc-toolset-11/root/usr/lib64/dyninst:/opt/rh/gcc-toolset-11/root/usr/lib/dyninst\n",
      "\u001b[2m\u001b[36m(_run pid=1826653)\u001b[0m 2023-09-26 20:05:08.016360: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/rh/gcc-toolset-11/root/usr/lib64:/opt/rh/gcc-toolset-11/root/usr/lib:/opt/rh/gcc-toolset-11/root/usr/lib64/dyninst:/opt/rh/gcc-toolset-11/root/usr/lib/dyninst\n",
      "\u001b[2m\u001b[36m(_run pid=1826653)\u001b[0m 2023-09-26 20:05:08.016462: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/rh/gcc-toolset-11/root/usr/lib64:/opt/rh/gcc-toolset-11/root/usr/lib:/opt/rh/gcc-toolset-11/root/usr/lib64/dyninst:/opt/rh/gcc-toolset-11/root/usr/lib/dyninst\n",
      "\u001b[2m\u001b[36m(_run pid=1826653)\u001b[0m 2023-09-26 20:05:08.016475: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.\n",
      " 26%|██▌       | 102/391 [00:46<01:59,  2.43it/s]\u001b[2m\u001b[36m(_run pid=1826989)\u001b[0m 2023-09-26 20:05:53.319477: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/rh/gcc-toolset-11/root/usr/lib64:/opt/rh/gcc-toolset-11/root/usr/lib:/opt/rh/gcc-toolset-11/root/usr/lib64/dyninst:/opt/rh/gcc-toolset-11/root/usr/lib/dyninst\n",
      "\u001b[2m\u001b[36m(_run pid=1826989)\u001b[0m 2023-09-26 20:05:54.012634: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/rh/gcc-toolset-11/root/usr/lib64:/opt/rh/gcc-toolset-11/root/usr/lib:/opt/rh/gcc-toolset-11/root/usr/lib64/dyninst:/opt/rh/gcc-toolset-11/root/usr/lib/dyninst\n",
      "\u001b[2m\u001b[36m(_run pid=1826989)\u001b[0m 2023-09-26 20:05:54.012721: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/rh/gcc-toolset-11/root/usr/lib64:/opt/rh/gcc-toolset-11/root/usr/lib:/opt/rh/gcc-toolset-11/root/usr/lib64/dyninst:/opt/rh/gcc-toolset-11/root/usr/lib/dyninst\n",
      "\u001b[2m\u001b[36m(_run pid=1826989)\u001b[0m 2023-09-26 20:05:54.012732: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.\n",
      "100%|██████████| 391/391 [02:53<00:00,  2.31it/s]\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866419)\u001b[0m /home/ssd2/zhaocaibei/miniconda3/envs/jupyter/lib/python3.8/site-packages/secretflow/ml/nn/sl/backend/torch/sl_base.py:627: UserWarning: The given NumPy array is not writable, and PyTorch does not support non-writable tensors. This means writing to this tensor will result in undefined behavior. You may want to copy the array to protect its data or make it writable before converting it to a tensor. This type of warning will be suppressed for the rest of this program. (Triggered internally at ../torch/csrc/utils/tensor_numpy.cpp:206.)\n",
      "\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866419)\u001b[0m   torch.as_tensor(e)\n",
      "100%|██████████| 391/391 [03:40<00:00,  1.78it/s, epoch: 1/1 -  train_loss:1.6854610443115234  train_MulticlassAccuracy:0.300819993019104  train_MulticlassPrecision:0.300819993019104  val_val_loss:1.873559594154358  val_MulticlassAccuracy:0.36000001430511475  val_MulticlassPrecision:0.36000001430511475 ]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866347)\u001b[0m Files already downloaded and verified\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866347)\u001b[0m INFO:root:Epoch: [1 | 1]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866347)\u001b[0m #Labeled: 40 #Unlabeled: 49960\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866347)\u001b[0m INFO:root:batch_idx: 0, loss: 2.3291807174682617\n",
      "\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866347)\u001b[0m INFO:root:batch_idx: 250, loss: 2.070810818102255\n",
      "\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866347)\u001b[0m INFO:root:batch_idx: 500, loss: 2.002920127676395\n",
      "\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866347)\u001b[0m INFO:root:batch_idx: 750, loss: 1.998070199226095\n",
      "\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866347)\u001b[0m INFO:root:batch_idx: 1000, loss: 2.016154585303841\n",
      "\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866347)\u001b[0m INFO:root:---Label inference on evaluation dataset 0\n",
      "\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866347)\u001b[0m INFO:root:evaluate 0\n",
      "\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866347)\u001b[0m INFO:root:evaluate 50\n",
      "\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866347)\u001b[0m INFO:root:Dataset Overall Statistics:\n",
      "\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866347)\u001b[0m INFO:root:top 1 accuracy:23.77, top 4 accuracy:66.64\n",
      "\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866347)\u001b[0m INFO:root:evaluate metric 0, 0.23770000040531158\n",
      "\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866347)\u001b[0m INFO:root:evaluate metric 1, 0.23770000040531158\n",
      "\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866347)\u001b[0m INFO:root:test_loss: 2.0509094161987305, test_acc: 23.77\n",
      "\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866347)\u001b[0m INFO:root:---Label inference on evaluation dataset 1\n",
      "\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866347)\u001b[0m INFO:root:evaluate 0\n",
      "\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866347)\u001b[0m INFO:root:evaluate 50\n",
      "\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866347)\u001b[0m INFO:root:evaluate 100\n",
      "\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866347)\u001b[0m INFO:root:evaluate 150\n",
      "\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866347)\u001b[0m INFO:root:evaluate 200\n",
      "\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866347)\u001b[0m INFO:root:evaluate 250\n",
      "\u001b[2m\u001b[36m(PYUSLTorchModel pid=1866347)\u001b[0m INFO:root:evaluate 300\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'train_loss': [array(1.685461, dtype=float32)], 'train_MulticlassAccuracy': [tensor(0.3008)], 'train_MulticlassPrecision': [tensor(0.3008)], 'val_val_loss': [array(1.8735596, dtype=float32)], 'val_MulticlassAccuracy': [tensor(0.3600)], 'val_MulticlassPrecision': [tensor(0.3600)]}\n"
     ]
    }
   ],
   "source": [
    "history = sl_model.fit(\n",
    "    fed_data,\n",
    "    label,\n",
    "    validation_data=(fed_data, label),\n",
    "    epochs=1,\n",
    "    batch_size=128,\n",
    "    shuffle=False,\n",
    "    random_seed=1234,\n",
    "    dataset_builder=None,\n",
    "    # callbacks=callback_dict, # 暂时注释掉，callback完成后恢复 @caibei\n",
    ")\n",
    "print(history)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bfeae819-6d79-4015-98c5-e657ba0649e5",
   "metadata": {},
   "source": [
    "## 总结\n",
    "本文通过 Cifar10 数据集上的标签推理攻击任务来演示如何通过隐语来使用 LabelInferenceAttack。\n",
    "\n",
    "您可以：\n",
    "\n",
    "1. 下载并拆分数据集，准备训练、攻击使用的数据\n",
    "2. 定义拆分模型结构及 SL Model\n",
    "3. 定义 attacker_builder，在其中定义攻击需要的 data_builder 和 LabelInfereceAttacker\n",
    "4. 调用 SL Model 进行训练攻击\n",
    "\n",
    "您可以在自己的数据集上进行尝试，如有任何问题，可以在 github 进行训练即可。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
